from .order_book import OrderBook
from decimal import Decimal

class OrderBookSZ(OrderBook):
    def __init__(self, pre_close_price:Decimal = Decimal("0.0")):
        self.pre_close_price: Decimal = pre_close_price
        super().__init__()

    def get_pre_close_price(self) -> Decimal:
        """
        获取前收盘价
        """
        return self.pre_close_price
    
    def set_pre_close_price(self, pre_close_price: Decimal):
        """
        设置前收盘价
        """
        self.pre_close_price = pre_close_price

    def get_auction_price(self) -> float:
        """ 计算集合竞价价格
        """
        all_prices = sorted(set(self._buy_orders.keys()) | set(self._sell_orders.keys()))
        price_to_volume = dict()
        max_volume = Decimal("0")

        # 计算每个价格的成交量
        for price in all_prices:
            buy_qty = sum(
                order.rest_quantity()
                for p, orders in self._buy_orders.items()
                if p >= price
                for order in orders
            )
            sell_qty = sum(
                order.rest_quantity()
                for p, orders in self._sell_orders.items()
                if p <= price
                for order in orders
            )
            match_volume = min(buy_qty, sell_qty)
            price_to_volume[price] = match_volume
            if match_volume > max_volume:
                max_volume = match_volume

        # 找出所有最大成交量对应的价格
        candidate_prices = [p for p, v in price_to_volume.items() if v == max_volume]

        if not candidate_prices:
            return None
        if len(candidate_prices) == 1:
            return candidate_prices[0]

        # 多个最大成交量价格，选离 pre_close_price 最近的那个
        best_price = min(candidate_prices, key=lambda p: abs(p - self.pre_close_price))
        return best_price